# visualizer.py
import math

import matplotlib.pyplot as plt

plt.rcParams['font.sans-serif'] = ['SimHei']  # 使用黑体
plt.rcParams['axes.unicode_minus'] = False  # 解决负号显示问题
from matplotlib.patches import Circle, RegularPolygon, Wedge
import numpy as np
from matplotlib.gridspec import GridSpec
import matplotlib.colors as mcolors
from matplotlib.lines import Line2D  # 导入Line2D用于绘制攻击连线

from agent import Drone, Boat


class Visualizer:
    def __init__(self, world, faction_to_display='white'):
        self.world = world
        self.faction_to_display = faction_to_display  # 要显示的阵营态势图
        self.fig = None
        self.ax_battle = None  # 战场态势图
        self.ax_situation = None  # 阵营态势图
        self.ax_obstacle = None  # 阵营障碍图
        self.visual_objects = {}
        self.text_objects = {}
        self.scan_objects = {}
        self.lock_objects = {}
        self.status_text = None
        self.last_draw_step = 0
        self.situation_title = None  # 态势图标题
        self.path_objects = {}  # 存储路径线对象
        self.waypoint_objects = {}  # 存储路径点对象
        self.attack_line_objects = {}  # 存储攻击连线对象
        self.obstacle_map_displayed = None  # 存储当前显示的障碍图
        self.defense_line = None  # 存储防线对象

    def init_visualization(self):
        """初始化可视化界面，分为3个部分"""
        plt.ion()

        # 创建图形和网格布局
        self.fig = plt.figure(figsize=(14, 8))
        gs = GridSpec(1, 3)

        # 左侧：战场态势图
        self.ax_battle = self.fig.add_subplot(gs[0])
        self.ax_battle.set_xlim(0, self.world.width)
        self.ax_battle.set_ylim(0, self.world.height)
        self.ax_battle.set_title('战场可视化')
        self.ax_battle.set_xlabel('X坐标 (米)')
        self.ax_battle.set_ylabel('Y坐标 (米)')
        self.ax_battle.grid(True)
        self.ax_battle.set_aspect('equal')
        # 中间：障碍图
        self.ax_obstacle = self.fig.add_subplot(gs[1])
        self.ax_battle.set_xlim(0, self.world.width)
        self.ax_battle.set_ylim(0, self.world.height)
        self.ax_obstacle.set_title(f"{'白方' if self.faction_to_display == 'white' else '黑方'}危险区分布")
        self.ax_obstacle.set_xlabel('X坐标 (米)')
        self.ax_obstacle.set_ylabel('Y坐标 (米)')
        self.ax_obstacle.grid(True)
        self.ax_battle.set_aspect('equal')
        # 右侧：阵营态势图
        self.ax_situation = self.fig.add_subplot(gs[2])
        self.situation_title = self.ax_situation.set_title(f"{'白方' if self.faction_to_display == 'white' else '黑方'}态势图")
        self.ax_situation.set_xlabel('X坐标 (米)')
        self.ax_situation.set_ylabel('Y坐标 (米)')
        self.ax_situation.grid(False)

        # 添加图例到战场态势图
        red_boat = plt.Line2D([0], [0], marker='>', markersize=8, color='r', linestyle='None')
        blue_boat = plt.Line2D([0], [0], marker='>', markersize=8, color='b', linestyle='None')
        red_drone = plt.Line2D([0], [0], marker='o', markersize=8, color='r', linestyle='None')
        sensor_range = Circle((0, 0), radius=100, color='gray', alpha=0.2, fill=True, linestyle='-')
        lock_range = Circle((0, 0), radius=100, color='gray', alpha=0.1, fill=False, linestyle='--')
        path_line = plt.Line2D([0, 1], [0, 1], color='gray', linestyle='--', linewidth=1)
        waypoint = plt.Line2D([0], [0], marker='o', markersize=5, color='gray', linestyle='None')
        attack_line = plt.Line2D([0, 1], [0, 1], color='gray', linestyle='-', linewidth=1)  # 添加攻击连线图例
        defense_line = plt.Line2D([0, 1], [0, 1], color='red', linestyle='-', linewidth=2)  # 添加防线图例项

        self.ax_battle.legend(
            [red_boat, blue_boat, red_drone, sensor_range, lock_range, path_line, waypoint, attack_line, defense_line],
            ['白方无人艇', '黑方无人艇', '白方无人机',
             '探测范围', '锁定范围', '规划航线', '航点', '攻击指示线', '防线'],
            loc='upper right',
            fontsize=7,
            framealpha=0.5)

        # 添加状态文本
        self.status_text = self.fig.text(
            0.02, 0.95,
            '',
            fontsize=10,
            bbox=dict(facecolor='white', alpha=0.7))
        self.update_status()

        # 添加控制说明
        controls_text = (
            "控制指令:\n"
            "+/-: 调整速度\n"
            "空格: 暂停/继续\n"
            "R: 切换范围显示\n"
            "T: 切换文本标签\n"
            "C: 截图\n"
            "Q: 退出"
        )
        self.fig.text(
            0.9, 0.9,
            controls_text,
            fontsize=9,
            ha='left',
            bbox=dict(facecolor='white', alpha=0.7)
        )

        # 连接键盘事件
        self.fig.canvas.mpl_connect('key_press_event', self.on_key_press)

        # 绘制防线
        self.draw_defense_line()
        # 初始化态势图显示
        self.update_situation_map()
        # 初始化障碍图显示
        self.update_obstacle_map()

    def on_key_press(self, event):
        """处理键盘事件，转发给World"""
        if hasattr(self, 'on_key_press_callback') and callable(self.on_key_press_callback):
            self.on_key_press_callback(event)

    def create_visual_object(self, agent):
        """为智能体创建可视化对象（在战场态势图上）"""
        color = 'red' if agent.faction == 'white' else 'blue'

        if isinstance(agent, Boat):
            # 船只用三角形表示
            boat = RegularPolygon(
                (agent.x, agent.y),
                3,
                radius=agent.radius * max(self.world.width, self.world.height) * 0.00015,
                orientation=agent.target_heading,
                color=color,
                alpha=1
            )
            self.visual_objects[agent.id] = self.ax_battle.add_patch(boat)

            # 添加文本标签
            text = self.ax_battle.text(
                agent.x,
                agent.y + agent.radius + 20,
                f'B{agent.id}\nHP:{agent.health}',
                fontsize=8,
                ha='center',
                va='center'
            )
            self.text_objects[agent.id] = text

            # 添加武器锁定范围（虚线圆）
            if agent.lock_range > 0:
                lock = Circle(
                    (agent.x, agent.y),
                    agent.lock_range,
                    color=color,
                    alpha=0.3,
                    fill=False,
                    linestyle='--',
                    linewidth=1,
                    visible=self.world.show_ranges
                )
                self.lock_objects[agent.id] = self.ax_battle.add_patch(lock)

            # 添加传感器范围（实心半透明圆）
            if agent.sensor_range > 0:
                scan = Circle(
                    (agent.x, agent.y),
                    agent.sensor_range,
                    color=color,
                    alpha=0.1,
                    fill=True,
                    linestyle='-',
                    linewidth=1,
                    visible=self.world.show_ranges
                )
                self.scan_objects[agent.id] = self.ax_battle.add_patch(scan)

        elif isinstance(agent, Drone):
            # 无人机用圆形表示
            drone = Circle(
                (agent.x, agent.y),
                radius=agent.radius * max(self.world.width, self.world.height) * 0.00005,
                color=color,
                alpha=1
            )
            self.visual_objects[agent.id] = self.ax_battle.add_patch(drone)

            # 添加文本标签
            text = self.ax_battle.text(
                agent.x,
                agent.y + agent.radius + 20,
                f'D{agent.id}\nP:{agent.power / 3600:.1f}h',
                fontsize=8,
                ha='center',
                va='center',
                visible=self.world.show_text
            )
            self.text_objects[agent.id] = text

            # 添加传感器范围（扇形，实心半透明）
            if agent.sensor_range > 0:
                # 计算扇形的起始和结束角度（转换为度）
                start_angle = np.degrees(agent.target_heading - agent.sensor_angle)
                end_angle = np.degrees(agent.target_heading + agent.sensor_angle)

                scan = Wedge(
                    (agent.x, agent.y),
                    agent.sensor_range,
                    start_angle,
                    end_angle,
                    color=color,
                    alpha=0 if agent.freeze else 0.1,
                    fill=True,
                    linestyle='-',
                    linewidth=1,
                    visible=self.world.show_ranges
                )
                self.scan_objects[agent.id] = self.ax_battle.add_patch(scan)

    def remove_visual_object(self, agent_id):
        """移除智能体的可视化对象"""
        if agent_id in self.visual_objects:
            self.visual_objects[agent_id].remove()
            del self.visual_objects[agent_id]
        if agent_id in self.text_objects:
            self.text_objects[agent_id].remove()
            del self.text_objects[agent_id]
        if agent_id in self.scan_objects:
            self.scan_objects[agent_id].remove()
            del self.scan_objects[agent_id]
        if agent_id in self.lock_objects:
            self.lock_objects[agent_id].remove()
            del self.lock_objects[agent_id]
        # 移除路径和航点
        if agent_id in self.path_objects:
            self.path_objects[agent_id].remove()
            del self.path_objects[agent_id]
        if agent_id in self.waypoint_objects:
            for obj in self.waypoint_objects[agent_id]:
                obj.remove()
            del self.waypoint_objects[agent_id]
        # 移除攻击连线
        if agent_id in self.attack_line_objects:
            self.attack_line_objects[agent_id].remove()
            del self.attack_line_objects[agent_id]

    def draw_defense_line(self):
        """绘制防线"""
        if hasattr(self.world, 'defense_line_x') and self.world.defense_line_x is not None:
            # 如果已存在防线对象，先移除
            if self.defense_line is not None:
                self.defense_line.remove()

            # 创建新的防线对象
            self.defense_line = self.ax_battle.axvline(
                x=self.world.defense_line_x,
                color='red',
                linewidth=2,
                alpha=0.7,
                linestyle='-'
            )

            # 添加防线标签
            self.ax_battle.text(
                self.world.defense_line_x - 15000,  # X位置偏移
                7500,  # Y位置在底部
                f'最终拦截线      最终拦截线        最终拦截线        最终拦截线        最终拦截线',
                fontsize=10,
                color='red',
                rotation=90,
                weight='bold'  # 加粗字体
            )

    def update_obstacle_map(self):
        """更新障碍图显示，并添加本阵营无人艇的位置和路径"""
        if self.ax_obstacle is None:
            return

        # 清除现有内容
        self.ax_obstacle.clear()

        # 设置障碍图的坐标轴范围与战场图一致
        self.ax_obstacle.set_xlim(0, self.world.width)
        self.ax_obstacle.set_ylim(0, self.world.height)
        self.ax_obstacle.set_xlabel('X坐标 (米)')
        self.ax_obstacle.set_ylabel('Y坐标 (米)')
        self.ax_obstacle.set_aspect('equal')

        # 检查指定阵营的态势图是否存在
        if self.faction_to_display in self.world.faction_maps:
            faction_map = self.world.faction_maps[self.faction_to_display]

            # 获取障碍图
            obstacle_map = faction_map.obstacle_map
            self.obstacle_map_displayed = obstacle_map

            # 绘制障碍图
            if obstacle_map is not None:
                # 创建自定义颜色映射：障碍区红色，安全区绿色
                cmap = mcolors.ListedColormap(['green', 'red'])

                # 计算障碍图的显示范围
                extent = [0, self.world.width, 0, self.world.height]

                # 绘制障碍图
                self.ax_obstacle.imshow(
                    obstacle_map.transpose(),  # 转置以匹配坐标方向
                    cmap=cmap,
                    origin='lower',
                    extent=extent,
                    alpha=0.3
                )

                # 添加图例说明
                safe_patch = plt.Rectangle((0, 0), 1, 1, fc='green', alpha=0.3)
                obstacle_patch = plt.Rectangle((0, 0), 1, 1, fc='red', alpha=0.3)
                self.ax_obstacle.legend(
                    [safe_patch, obstacle_patch],
                    ['安全区域', '危险区域'],
                    loc='upper left'
                )

                # 添加本阵营无人艇的位置和路径
                for agent in self.world.agents:
                    if agent.faction == self.faction_to_display and isinstance(agent, Boat):
                        # 绘制无人艇位置
                        boat = RegularPolygon(
                            (agent.x, agent.y),
                            3,
                            radius=agent.radius * max(self.world.width, self.world.height) * 0.00015,
                            orientation=agent.target_heading + math.pi / 6,
                            color='red',
                            alpha=1
                        )
                        self.ax_obstacle.add_patch(boat)

                        # 添加文本标签
                        self.ax_obstacle.text(
                            agent.x,
                            agent.y + agent.radius + 20,
                            f'无人艇{agent.id}\n生命:{agent.health}\n速度:{agent.v:.1f}',
                            fontsize=8,
                            ha='center',
                            va='center'
                        )

                        # 绘制路径（如果有）
                        if hasattr(agent, 'controller') and hasattr(agent.controller,
                                                                    'waypoints') and agent.controller.waypoints:
                            waypoints = agent.controller.waypoints
                            path_points = [(agent.x, agent.y)]

                            for wp in waypoints:
                                position = wp.get('position', [])
                                if len(position) >= 2:
                                    path_points.append((position[0], position[1]))

                            if len(path_points) > 1:
                                path_x, path_y = zip(*path_points)
                                self.ax_obstacle.plot(
                                    path_x, path_y,
                                    color='red',
                                    linestyle='--',
                                    linewidth=1.5,
                                    alpha=0.7,
                                    marker='o',
                                    markersize=4
                                )

                # 设置标题
                self.ax_obstacle.set_title(
                    f"{'白方' if self.faction_to_display == 'white' else '黑方'}危险区分布")
            else:
                self.ax_obstacle.text(0.5, 0.5, "无可用障碍图",
                                      ha='center', va='center', fontsize=12)
        else:
            self.ax_obstacle.text(0.5, 0.5,
                                  f"无{'白方' if self.faction_to_display == 'white' else '黑方'}态势图",
                                  ha='center', va='center', fontsize=12)

        # 添加网格线
        self.ax_obstacle.grid(True, linestyle='--', alpha=0.3)

    def update_situation_map(self):
        """更新阵营态势图显示"""
        if self.ax_situation is None:
            return

        # 清除现有内容
        self.ax_situation.clear()
        self.ax_situation.set_xlabel('X坐标 (米)')
        self.ax_situation.set_ylabel('Y坐标 (米)')

        # 检查指定阵营的态势图是否存在
        if self.faction_to_display in self.world.faction_maps:
            # 获取并绘制态势图
            faction_map = self.world.faction_maps[self.faction_to_display]
            faction_map.plot(self.ax_situation,
                             title=f"{'白方' if self.faction_to_display == 'white' else '黑方'}态势图",
                             cmap='hot')  # 使用热力图颜色映射

            # 添加网格线
            self.ax_situation.grid(True, linestyle='--', alpha=0.3)

            # 添加比例尺
            self.ax_situation.text(0.02, 0.02,
                                   f"网格分辨率: {faction_map.grid_resolution / 1000:.1f}km",
                                   transform=self.ax_situation.transAxes,
                                   fontsize=8,
                                   bbox=dict(facecolor='white', alpha=0.7))
        else:
            # 如果态势图不存在，显示提示信息
            self.ax_situation.text(0.5, 0.5,
                                   f"无{'白方' if self.faction_to_display == 'white' else '黑方'}态势图",
                                   ha='center', va='center',
                                   fontsize=12)
            self.ax_situation.set_title(f"{'白方' if self.faction_to_display == 'white' else '黑方'}态势图")
            self.ax_situation.grid(False)

    def update_visualization(self):
        """更新整个可视化界面，包括战场态势和阵营态势图"""
        # 更新战场态势图
        for agent in self.world.agents:
            if agent.id in self.visual_objects:
                # 更新位置和方向
                if isinstance(agent, Boat):
                    self.visual_objects[agent.id].xy = (agent.x, agent.y)
                    self.visual_objects[agent.id].orientation = agent.target_heading + math.pi / 6
                elif isinstance(agent, Drone):
                    self.visual_objects[agent.id].center = (agent.x, agent.y)

                # 更新文本
                if isinstance(agent, Boat):
                    self.text_objects[agent.id].set_text(
                        f'无人艇{agent.id}' if agent.faction is self.faction_to_display else f'无人艇{agent.id}\n生命:{agent.health}\n速度:{agent.v:.1f}')
                    self.text_objects[agent.id].set_position((agent.x, agent.y + agent.radius + 20))
                    self.text_objects[agent.id].set_visible(self.world.show_text)
                elif isinstance(agent, Drone):
                    self.text_objects[agent.id].set_text(
                        f'无人机{agent.id}\n电量:{agent.power:.1f}s\n速度:{agent.v:.1f}')
                    self.text_objects[agent.id].set_position((agent.x, agent.y + agent.radius + 20))
                    self.text_objects[agent.id].set_visible(self.world.show_text)

                # 更新武器锁定范围（Boat专用）
                if agent.id in self.lock_objects and isinstance(agent, Boat):
                    self.lock_objects[agent.id].center = (agent.x, agent.y)
                    self.lock_objects[agent.id].set_visible(self.world.show_ranges)

                # 更新传感器范围
                if agent.id in self.scan_objects:
                    if isinstance(agent, Boat):
                        self.scan_objects[agent.id].center = (agent.x, agent.y)
                        self.scan_objects[agent.id].set_visible(self.world.show_ranges)
                    elif isinstance(agent, Drone):
                        # 移除旧的扫描范围对象
                        self.scan_objects[agent.id].remove()
                        # 创建新的扫描范围对象
                        color = 'red' if agent.faction == 'white' else 'blue'
                        # 计算扇形的起始和结束角度（转换为度）
                        start_angle = np.degrees(math.atan2(agent.vy, agent.vx) - agent.sensor_angle)
                        end_angle = np.degrees(math.atan2(agent.vy, agent.vx) + agent.sensor_angle)

                        scan = Wedge(
                            (agent.x, agent.y),
                            agent.sensor_range,
                            start_angle,
                            end_angle,
                            color=color,
                            alpha=0 if agent.freeze else 0.1,
                            fill=True,
                            linestyle='-',
                            linewidth=1,
                            visible=self.world.show_ranges
                        )
                        self.scan_objects[agent.id] = self.ax_battle.add_patch(scan)

                # 更新路径和航点
                self.update_path_visualization(agent)

        # 添加攻击连线
        self.update_attack_lines()

        # 更新战场可视化标题
        battle_title = f'战场可视化 - 时间: {self.world.time:.1f}秒'
        if self.world.paused:
            battle_title += " [已暂停]"
        self.ax_battle.set_title(battle_title)

        # 更新障碍图
        self.update_obstacle_map()

        # 更新阵营态势图
        self.update_situation_map()

        # 更新状态文本
        self.update_status()

        # 重绘
        self.fig.canvas.draw()
        self.fig.canvas.flush_events()

    def update_attack_lines(self):
        """更新攻击连线显示"""
        # 先移除所有现有的攻击连线
        for line in list(self.attack_line_objects.values()):
            line.remove()
        self.attack_line_objects.clear()

        # 遍历所有Boat智能体
        for agent in self.world.agents:
            if isinstance(agent, Boat) and agent._target is not None and not agent.freeze:
                # 获取攻击者和目标的位置
                attacker_pos = (agent.x, agent.y)
                target = agent._target

                # 如果目标已被移除，则跳过
                if target not in self.world.agents or agent.distance_to(target) > agent.lock_range:
                    continue

                target_pos = (target.x, target.y)

                # 根据攻击者的阵营设置连线颜色
                line_color = 'red' if agent.faction == 'white' else 'blue'

                # 创建攻击连线
                attack_line = Line2D(
                    [attacker_pos[0], target_pos[0]],
                    [attacker_pos[1], target_pos[1]],
                    color=line_color,
                    linewidth=1.5,
                    linestyle='-',
                    alpha=0.7
                )

                # 添加到图表和字典中
                self.attack_line_objects[agent.id] = self.ax_battle.add_line(attack_line)

    def update_path_visualization(self, agent):
        """更新智能体的路径可视化"""
        agent_id = agent.id

        # 检查控制器和路径点是否存在
        if hasattr(agent, 'controller') and hasattr(agent.controller, 'waypoints') and agent.controller.waypoints:
            waypoints = agent.controller.waypoints
            # 创建路径点坐标列表（从当前位置开始）
            path_points = [(agent.x, agent.y)]
            for wp in waypoints:
                position = wp.get('position', [])
                if len(position) >= 2:  # 确保有x和y坐标
                    path_points.append((position[0], position[1]))

            if len(path_points) > 1:
                # 转换坐标列表为x和y的单独列表
                path_x, path_y = zip(*path_points)

                # 更新或创建路径线
                if agent_id in self.path_objects:
                    self.path_objects[agent_id].set_data(path_x, path_y)
                else:
                    path_line, = self.ax_battle.plot(
                        path_x, path_y,
                        color='red' if agent.faction == 'white' else 'blue',
                        linestyle='--',
                        linewidth=0.5,
                        alpha=0.5,
                        marker=None
                    )
                    self.path_objects[agent_id] = path_line

                # 更新或创建航点标记
                waypoint_markers = []
                if agent_id in self.waypoint_objects:
                    for obj in self.waypoint_objects[agent_id]:
                        obj.remove()

                # 跳过第一个点（当前位置）
                for x, y in path_points[1:]:
                    marker = self.ax_battle.plot(
                        x, y,
                        marker='o',
                        markersize=3,
                        color='red' if agent.faction == 'white' else 'blue',
                        alpha=0.5,
                        linestyle='None'
                    )[0]
                    waypoint_markers.append(marker)

                self.waypoint_objects[agent_id] = waypoint_markers
            else:
                # 路径点不足，移除现有路径
                self.remove_path(agent_id)
        else:
            # 没有路径点，移除现有路径
            self.remove_path(agent_id)

    def remove_path(self, agent_id):
        """移除智能体的路径可视化"""
        if agent_id in self.path_objects:
            self.path_objects[agent_id].remove()
            del self.path_objects[agent_id]
        if agent_id in self.waypoint_objects:
            for obj in self.waypoint_objects[agent_id]:
                obj.remove()
            del self.waypoint_objects[agent_id]

    def update_status(self):
        """更新状态文本"""
        if self.status_text:
            self.status_text.set_text(
                f'速度: {self.world.speed_factor:.1f}x | 暂停: {self.world.paused} | '
                f'范围显示: {"ON" if self.world.show_ranges else "OFF"} | '
                f'文字显示: {"ON" if self.world.show_text else "OFF"}'
            )

    def pause_handler(self):
        """处理暂停状态"""
        plt.pause(0.1)

    def is_active(self):
        """检查可视化窗口是否仍然活跃"""
        return plt.fignum_exists(self.fig.number) if self.fig else False

    def show(self):
        """显示可视化窗口"""
        plt.ioff()
        plt.show()

    def close(self):
        """关闭可视化窗口"""
        if self.fig:
            plt.close(self.fig)
